#pragma once

#include "physical_operator_type.hpp"
#include "value_type.hpp"
#include <arrow/api.h>
#include "expression_type.hpp"
#include <functional>

namespace DaseX {

// 记录该行数据所在位置
struct Offset {
    int64_t chunk_id;
    int64_t offset;
    Offset(int64_t chunk_id, int64_t offset) : chunk_id(chunk_id), offset(offset) {}
};

using KEY = Value;
using VALUE = Offset;
// 桶存储的元素，key + value,key是JOIN的键，Value是实际的一行数据，这里我们用Offset来表示一行数据
struct EntrySingle {
    std::vector<std::shared_ptr<KEY>> keys;
    VALUE val;
    size_t hash_val;
    EntrySingle(std::vector<std::shared_ptr<KEY>> keys, VALUE val, size_t hash_val) : keys(std::move(keys)), val(val), hash_val(hash_val) {}
};

// 行数据分区，也可以理解为HashTable的桶，用来存储数据
struct TupleBucket {
    int tuple_nums = 0;
    std::vector<std::shared_ptr<EntrySingle>> entry_set;
    void InsertEntry(std::shared_ptr<EntrySingle> &entrySingle) {
        entry_set.emplace_back(entrySingle);
        tuple_nums++;
    }
    std::shared_ptr<EntrySingle> GetEntry(int idx) {
        return entry_set[idx];
    }
};
// 记录探测状态
struct ProbeState {
    std::vector<int> match;
    // TODO: 可用用更节省内存的方式表示bit_map，后续优化
    std::vector<int> bit_map;
    std::unordered_map<int, std::vector<Offset>> match_result;
    void AddMatch(int idx) {
        match.emplace_back(idx);
    }
    void AddMatchResult(int idx, std::vector<Offset> &res) {
        match_result[idx] = res;
    }
};

class JoinHashTable {
public:
    JoinTypes join_type;
    std::vector<int> build_ids;
    std::vector<int> probe_ids;
    // TODO: 多比较条件暂时只用于SEMI、ANTI-join，应该扩展到所有类型JOIN，后续优化
    std::vector<ExpressionTypes> comparison_types;
    std::vector<std::shared_ptr<TupleBucket>> buckets;
    std::vector<std::shared_ptr<arrow::RecordBatch>> data_chunks;
    // 用于外连接补空
    std::shared_ptr<arrow::Schema> schema;
    std::shared_ptr<arrow::RecordBatch> null_row;
    int64_t nums_chunk = 0;
    const uint64_t radix_bits = 16;
    uint64_t NUM_PARTITIONS;
    uint64_t SHIFT;
    uint64_t MASK;
    int64_t total_count = 0;
    int precision = 2;
public:
    JoinHashTable(std::vector<int> build_ids, std::vector<int> probe_ids, std::vector<ExpressionTypes> comparison_types, JoinTypes join_type = JoinTypes::INNER);
    void Build(std::shared_ptr<arrow::RecordBatch> &chunk);
    std::shared_ptr<ProbeState> Probe(std::shared_ptr<arrow::RecordBatch> &chunk);
    // 根据Probe结果，获取最终的连接数据
    template<JoinTypes join_type>
    void GetJoinResult(std::shared_ptr<arrow::RecordBatch> &chunk, std::shared_ptr<arrow::RecordBatch> &result, std::shared_ptr<ProbeState> &probe_state, std::shared_ptr<ProbeState> &last_probe_state, bool is_final_probe = false);
public:
    // 对输入数组求Hash
    std::vector<size_t> Hashs(std::vector<std::shared_ptr<arrow::Array>> &join_keys, int count);
    // 对输入Hash数组求每条数据对应的桶编号
    std::vector<int> ComputeBucketIndices(std::vector<size_t> &join_key_hash, int count);
    // 根据桶编号将数据散列到各个桶
    void ScatterData(std::vector<std::shared_ptr<arrow::Array>> &join_keys, std::vector<int> &bucket_idx, int count, std::vector<size_t> &hashs);
    // 根据桶编号匹配所有数据
    std::shared_ptr<ProbeState> GatherData(std::vector<std::shared_ptr<arrow::Array>> &join_keys, std::vector<int> &bucket_idx, int count, std::vector<size_t> &hashes);
    // 用于SEMI join
    std::shared_ptr<ProbeState> GatherSemiData(std::vector<std::shared_ptr<arrow::Array>> &join_keys,
                                                              std::vector<int> &bucket_idx,
                                                              int count,
                                                              std::vector<size_t> &hashes);
    // 用于ANTI join
    std::shared_ptr<ProbeState> GatherAntiData(std::vector<std::shared_ptr<arrow::Array>> &join_keys,
                                               std::vector<int> &bucket_idx,
                                               int count,
                                               std::vector<size_t> &hashes);
    std::shared_ptr<arrow::RecordBatch> GetSingleRow(Offset offset);
    // 构建一行空数据
    void GetSingleNullRow();
    // Radix分区相关方法
    static inline constexpr uint64_t NumberOfPartitions(uint64_t radix_bits) {
        return uint64_t(1) << radix_bits;
    }
    //! Radix bits begin after uint16_t because these bits are used as salt in the aggregate HT
    static inline constexpr uint64_t Shift(uint64_t radix_bits) {
        return (sizeof(uint64_t) - sizeof(uint16_t)) * 8 - radix_bits;
    }
    //! Mask of the radix bits of the hash
    static inline constexpr uint64_t Mask(uint64_t radix_bits) {
        return (uint64_t(1 << radix_bits) - 1) << Shift(radix_bits);
    }
};

} // DaseX
